import os
import shutil
import random
from tqdm import tqdm  # 用于显示进度条


def split_yolo_dataset(dataset_path, train_ratio=0.8):
    """
    分割YOLO格式数据集为训练集和验证集

    参数:
    dataset_path -- 数据集根目录路径
    train_ratio -- 训练集比例 (默认0.8)
    """
    # 定义原始路径
    images_dir = os.path.join(dataset_path, 'images')
    labels_dir = os.path.join(dataset_path, 'labels')

    # 创建目标文件夹 - 使用用户要求的目录结构
    train_img_dir = os.path.join(dataset_path, 'images', 'train')
    val_img_dir = os.path.join(dataset_path, 'images', 'val')
    train_lbl_dir = os.path.join(dataset_path, 'labels', 'train')
    val_lbl_dir = os.path.join(dataset_path, 'labels', 'val')

    # 一次性创建所有需要的目录
    for path in [train_img_dir, val_img_dir, train_lbl_dir, val_lbl_dir]:
        os.makedirs(path, exist_ok=True)

    # 获取所有图像文件名 (不带扩展名)
    image_files = []
    for f in os.listdir(images_dir):
        if f.lower().endswith(('.png', '.jpg', '.jpeg')):
            # 跳过已存在的train/val子目录中的文件
            if f == 'train' or f == 'val':
                continue
            image_files.append(os.path.splitext(f)[0])

    # 随机打乱文件列表
    random.shuffle(image_files)

    # 计算分割点
    split_index = int(len(image_files) * train_ratio)

    # 分割文件列表
    train_files = image_files[:split_index]
    val_files = image_files[split_index:]

    print(f"总样本数: {len(image_files)}")
    print(f"训练集数量: {len(train_files)}")
    print(f"验证集数量: {len(val_files)}")

    # 支持的图像扩展名列表
    image_extensions = ['.png', '.jpg', '.jpeg']

    # 复制训练集文件
    print("\n正在处理训练集...")
    for file in tqdm(train_files):
        # 复制图像
        copied = False
        for ext in image_extensions:
            src_img = os.path.join(images_dir, file + ext)
            if os.path.exists(src_img):
                shutil.copy(src_img, train_img_dir)
                copied = True
                break

        # 复制标签
        src_lbl = os.path.join(labels_dir, file + '.txt')
        if os.path.exists(src_lbl):
            shutil.copy(src_lbl, train_lbl_dir)
        elif copied:
            print(f"\n警告: 找不到训练集标签文件 {file}.txt")

    # 复制验证集文件
    print("\n正在处理验证集...")
    for file in tqdm(val_files):
        # 复制图像
        copied = False
        for ext in image_extensions:
            src_img = os.path.join(images_dir, file + ext)
            if os.path.exists(src_img):
                shutil.copy(src_img, val_img_dir)
                copied = True
                break

        # 复制标签
        src_lbl = os.path.join(labels_dir, file + '.txt')
        if os.path.exists(src_lbl):
            shutil.copy(src_lbl, val_lbl_dir)
        elif copied:
            print(f"\n警告: 找不到验证集标签文件 {file}.txt")

    print("\n数据集分割完成！")
    print(f"训练集图像保存至: {train_img_dir}")
    print(f"训练集标签保存至: {train_lbl_dir}")
    print(f"验证集图像保存至: {val_img_dir}")
    print(f"验证集标签保存至: {val_lbl_dir}")


# 使用方法
if __name__ == "__main__":
    # 设置你的数据集路径 (包含images和labels文件夹的目录)
    dataset_path = r"D:\Datasets\pic0910-公司未戴头盔软件"

    split_yolo_dataset(dataset_path, train_ratio=0.8)